#!/usr/bin/python
# -*- coding: utf-8 -*-

'''
 PyIPS - Python image processing server/script

 author: Sergey Polzunov aka Traut
 email: traut.box[at]gmail.com
 blog: http://out.com.ua/blog

 Version 0.1
'''

import string

chunk_min_length = 4 + 4 + 4 # data_length, chunk_name, crc

main_chunks = ["IHDR", "PLTE", "IDAT", "IEND"]
other_chunks = ["tRNS",					# transparency
		"cHRM", "gAMA", "iCCP", "sBIT", "sRGB",	# colorspace
		"iTXt", "tEXt", "zTXt", "bKGD",		# textual info
		"hIST", "pHYs", "sPLT",			# miscellaneous things
		"tIME"]					# time stamp

def compute_chunk_length(data, i):
    return (ord(data[i]) << 3) | (ord(data[i + 1]) << 2) | (ord(data[i + 2]) << 1) | ord(data[i + 3])

def remove_chunk(data, chunk_name):
    for i in range(0, len(data) - chunk_min_length + 1):
	j = i + 4
	maybe_chunk = data[j] + data[j + 1] + data[j + 2] + data[j + 3]
	if maybe_chunk == chunk_name and (main_chunks + other_chunks).count(maybe_chunk) == 1:
	    chunk_data_length = compute_chunk_length(data, i);
	    #print "removing %d blocks" % (chunk_min_length + chunk_data_length)
	    return data[:i] + data[i + chunk_min_length + chunk_data_length:]
    print "Cannot find chunk with name", chunk_name


def browse_chunks(data):
    found_chunks = []
    for i in range(0, len(data) - chunk_min_length + 1):
	j = i + 4
	maybe_chunk = data[j] + data[j + 1] + data[j + 2] + data[j + 3]
	if (main_chunks + other_chunks).count(maybe_chunk) == 1:
	    chunk_data_length = compute_chunk_length(data, i);
	    found_chunks.append((maybe_chunk, chunk_data_length))
	    i += chunk_data_length + chunk_min_length
    
    return found_chunks
    
    #for chunk_name, chunk_length in found_chunks:
    #	print "Chunk %s founded. Data length is %d" % (chunk_name, chunk_length)
	
	
    
def test():
    import Image
    #image = Image.open("")
    image_file = open("test.png", "rb")
    image_raw_data = image_file.read()
    browse_chunks(image_raw_data)
    chunk = "gAMA"
    print "Removing", chunk
    new_image_data = remove_chunk(image_raw_data, chunk)
    new_image_file = open("test2.png", "wb")
    new_image_file.write(new_image_data)
    
    image_file.close()
    new_image_file.close()
    
    image_file = open("test2.png", "rb")
    image_raw_data = image_file.read()
    browse_chunks(image_raw_data)
    image_file.close()

if __name__ == "__main__":
    test()
    
